import pyrealsense2 as rs
import cv2
import numpy as np

def capture_and_save(
    serial_number,
    rgb_path="rgb_image.png",
    depth_rgb_path="depth_rgb.png",  # 彩色可视化的深度图
    depth_npy_path="depth_data.npy"  # 原始深度数据（npy格式）
):
    # 1. 配置相机管道
    pipeline = rs.pipeline()
    config = rs.config()

    # 2. 筛选指定序列号的设备
    ctx = rs.context()
    devices = ctx.query_devices()
    target_device = None
    for dev in devices:
        if dev.get_info(rs.camera_info.serial_number) == serial_number:
            target_device = dev
            break
    if not target_device:
        raise ValueError(f"未找到序列号为 {serial_number} 的设备，请检查连接或序列号")

    # 3. 配置流参数（RGB和深度，分辨率1280x720，帧率30）
    # 注意：若需"720x1280"竖屏，将宽高参数改为 (720, 1280)，但需确认相机支持
    config.enable_device(serial_number)
    config.enable_stream(rs.stream.color, 1280, 720, rs.format.bgr8, 30)  # RGB流（BGR格式）
    config.enable_stream(rs.stream.depth, 1280, 720, rs.format.z16, 30)   # 深度流（16位原始数据）

    # 4. 启动管道并获取帧
    pipeline.start(config)
    try:
        # 跳过前10帧，等待相机稳定
        for _ in range(10):
            pipeline.wait_for_frames()

        # 获取一帧数据
        frames = pipeline.wait_for_frames()
        color_frame = frames.get_color_frame()
        depth_frame = frames.get_depth_frame()

        if not color_frame or not depth_frame:
            raise RuntimeError("无法获取RGB或深度帧，请检查相机配置")

        # 5. 转换为数组格式
        # RGB图：BGR格式（OpenCV兼容）
        color_image = np.asanyarray(color_frame.get_data())
        # 深度图：原始16位数据（单位：毫米）
        depth_raw = np.asanyarray(depth_frame.get_data())  # 这是需要保存为npy的原始数据
        # 深度图可视化（彩色映射，便于查看）
        depth_colormap = cv2.applyColorMap(
            cv2.convertScaleAbs(depth_raw, alpha=0.03),  # alpha调整对比度
            cv2.COLORMAP_JET
        )

        # 6. 保存数据
        # 保存RGB图
        cv2.imwrite(rgb_path, color_image)
        # 保存彩色可视化深度图（RGB格式）
        cv2.imwrite(depth_rgb_path, depth_colormap)
        # 保存原始深度数据（npy格式，保留16位数值）
        np.save(depth_npy_path, depth_raw)

        print(f"RGB图已保存至：{rgb_path}")
        print(f"彩色深度图已保存至：{depth_rgb_path}")
        print(f"原始深度数据（npy）已保存至：{depth_npy_path}")

    finally:
        pipeline.stop()  # 释放相机资源

if __name__ == "__main__":
    camera_serial = "318122301415"  # 相机序列号
    # 调用函数，可自定义保存路径
    capture_and_save(
        camera_serial,
        rgb_path="inputs/input_rgb_2.png",
        depth_rgb_path="inputs/input_depth_rgb_2.png",
        depth_npy_path="inputs/input_depth_raw_2.npy"
    )
